#pragma once

#include <stdint.h>
#include <unordered_set>

#include <google/protobuf/repeated_field.h>

#include "index/primary_key_index/primary_key_indexer.h"
#include "index/deletion_index/deletion_index_indexer.h"
#include "file_system/file_reader/mmap_file_reader.h"

namespace hawking {
namespace indexlib {

class PrimaryKeyDiskIndexer : public PrimaryKeyIndexer {
public:
    static const uint32_t pk_length_ = sizeof(uint64_t);
    static const uint32_t docid_length_ = sizeof(DocId);
    static const uint32_t pk_disk_item_length_ = pk_length_ + docid_length_;

    PrimaryKeyDiskIndexer() = default;
    virtual ~PrimaryKeyDiskIndexer() = default;

    // arg2 is deleted pks when dumping or merging
    Status Load(
        const std::string& pk_dir,
        std::shared_ptr<DeletionMapIndexer> deletion_indexer);

    virtual Status Lookup(
        const google::protobuf::RepeatedField<int64_t>& pks,
        SegmentContext* segment_ctx) override;
    virtual ReturnValue<DocId> Lookup(uint64_t pk) override;

    bool IsExist(uint64_t pk) const;

    virtual uint32_t DocCount() const override {
        return (pk_disk_reader_->Length() / pk_disk_item_length_);
    }

    void ExcludeTombPks(
        const HashSet<uint64_t>& tomb_stone_pks,
        std::shared_ptr<DeletionMapIndexer> deletion_indexer);

    const MmapFileReader& FileReader() const {
        return *pk_disk_reader_;
    }

    class PrimaryKeyDiskIterator : public PrimaryKeyIterator {
    public:
        PrimaryKeyDiskIterator(const MmapFileReader& pk_disk_reader) :
            pk_disk_reader_(&pk_disk_reader) {}
        virtual ~PrimaryKeyDiskIterator() = default;

        virtual bool HasNext() const override {
            return pk_disk_reader_ && current_offset_ < pk_disk_reader_->Length();
        }

        virtual void Next() {
            current_offset_ += sizeof(uint64_t);
            current_offset_ += sizeof(DocId);
        }

        virtual void Value(uint64_t* pk, DocId* doc_id) const {
            pk_disk_reader_->Read((uint8_t*)pk, sizeof(*pk), current_offset_);
            pk_disk_reader_->Read(
                (uint8_t*)doc_id, sizeof(*doc_id), current_offset_ + sizeof(*pk));
        }

        std::pair<uint64_t, DocId> Value() const {
            uint64_t pk;
            DocId doc_id;
            pk_disk_reader_->Read((uint8_t*)&pk, sizeof(pk), prev_offset);
            pk_disk_reader_->Read(
                (uint8_t*)&doc_id, sizeof(doc_id), prev_offset + sizeof(pk));
            return std::make_pair(pk, doc_id);
        }

    private:
        const MmapFileReader* pk_disk_reader_ = nullptr;
        size_t current_offset_ = 0;
    };

    virtual std::shared_ptr<PrimaryKeyIterator> CreateIterator() const override {
        return std::make_shared<PrimaryKeyDiskIterator>(FileReader());
    }

private:
    void LookupOne_(uint64_t pk, SegmentContext* segment_ctx);

    std::unique_ptr<MmapFileReader> pk_disk_reader_;
    std::shared_ptr<MMapAllocator> allocator_;

    uint64_t min_pk_ = 0;
    uint64_t max_pk_ = 0;
};

}
}